"""
如何使用上下文压缩进行检索?
检索的一个挑战是，通常在将数据导入系统时，你并不知道文档存储系统将面临的具体查询。这意味着与查询最相关的信息可能会埋藏在包含大量无关文本的文档中。
将整个文档传递给你的应用程序可能会导致更昂贵的LLM调用和更差的响应。

上下文压缩旨在解决这个问题。这个想法很简单：你可以使用给定查询的上下文来压缩检索到的文档，而不是立即按原样返回它们，以便只返回相关信息。
这里的“压缩”既指压缩单个文档的内容，也指整体过滤掉文档。
"""
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor, LLMChainFilter, LLMListwiseRerank, \
    EmbeddingsFilter
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import CharacterTextSplitter
from models import get_ollama_embeddings_client, get_ds_model_client

def pretty_print_docs(docs):
    print(
        f"\n{'-' * 100}\n".join(
            [f"Document {i+1}:\n\n" + d.page_content for i, d in enumerate(docs)]
        )
    )

file_path = "../data/document/state_of_the_union.txt"
documents = TextLoader(file_path, encoding="utf-8").load()
splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = splitter.split_documents(documents)

embeddings_client = get_ollama_embeddings_client()
retriever = Chroma.from_documents(texts, embedding=embeddings_client).as_retriever()

# docs = retriever.invoke("What did the president say about Ketanji Brown Jackson")
# print(docs)

llm_client = get_ds_model_client()

# 使用 LLMChainExtractor 添加上下文压缩
compressor = LLMChainExtractor.from_llm(llm_client)
# 其他压缩器：LLMChainFilter是一个稍微简单但更强大的压缩器，它使用 LLM 链来决定过滤掉哪些最初检索到的文档，以及返回哪些文档，而不操纵文档内容。
# compressor = LLMChainFilter.from_llm(llm_client)

compression_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)
compress_docs = compression_retriever.invoke("What did the president say about Ketanji Jackson Brown")
print(pretty_print_docs(compress_docs))



